# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team
import torch
from torch import nn


## Note that the following code is modified from
## https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/reward_model.py
class RewardModel(nn.Module):

    def __init__(self,
                 base_model,
                 tokenizer,
                 num_padding_at_beginning=0,
                 compute_fp32_loss=False):
        super().__init__()
        # 基础模型的配置信息
        self.config = base_model.config
        # 序列开始填充数量
        self.num_padding_at_beginning = num_padding_at_beginning
        # 检查模型是否包含word_embed_proj_dim属性
        if hasattr(self.config, "word_embed_proj_dim"):
            # `OPT` models use word_embed_proj_dim as final output
            # https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py#L497
            # 使用一个线性层self.v_head将word_embed_proj_dim映射到1，这个线性层没有偏置
            self.v_head = nn.Linear(self.config.word_embed_proj_dim,
                                    1,
                                    bias=False)
        else:
            # `gpt-neo(x)` models use `hidden_size` attribute names instead of `n_embd``
            # 检查配置是否包含hidden_size属性。如果包含，那么将其赋值给n_embd
            self.config.n_embd = self.config.hidden_size if hasattr(
                self.config, "hidden_size") else self.config.n_embd
            # 使用一个线性层self.v_head将word_embed_proj_dim映射到1，这个线性层没有偏置
            self.v_head = nn.Linear(self.config.n_embd, 1, bias=False)
        # 基础模型
        self.rwtransformer = base_model
        # 填充token ID
        self.PAD_ID = tokenizer.pad_token_id
        # 是否使用单精度计算损失函数
        self.compute_fp32_loss = compute_fp32_loss
    # 启用梯度检查点，减少内存使用
    def gradient_checkpointing_enable(self):
        self.rwtransformer.gradient_checkpointing_enable()
    # 禁用梯度检查点
    def gradient_checkpointing_disable(self):
        self.rwtransformer.gradient_checkpointing_disable()

    def forward(self,
                input_ids=None, # 模型输入数据
                past_key_values=None, # 这个参数用于Transformer模型中的deoder。在进行序列生成任务时，decoder可以使用之前时间不步的隐藏状态    
                attention_mask=None, # 指定哪些元素不应该被注意力机制关注到
                position_ids=None, # 每个输入token的位置编码
                head_mask=None, # 用于屏蔽某些注意力头
                inputs_embeds=None, # 非None时，模型将使用此作为输入而不是input_ids
                use_cache=False # 是否使用缓存来加速解码
        ):
        loss = None

        if self.config.model_type == "llama":
            kwargs = dict()
        else:
            kwargs = dict(head_mask=head_mask)

        transformer_outputs = self.rwtransformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            **kwargs)
        # transformer 模型的主要输出，表示模型对输入的隐藏表示
        hidden_states = transformer_outputs[0]
        # 将hidden_states传递到self.v_head(一个全连接层)并删除最后一个维度(通过squeeze(-1))来计算reward
        # 这个值可以被视为模型对每个输入token的奖励
        rewards = self.v_head(hidden_states).squeeze(-1)
        # 模型选择或拒绝某些特定输入的平均数
        chosen_mean_scores = []
        rejected_mean_scores = []

        # Split the inputs and rewards into two parts, chosen and rejected
        # input_ids的形状是[batch_size, sequence_length]
        assert len(input_ids.shape) == 2
        # 批处理大小的一半
        bs = input_ids.shape[0] // 2
        # 序列长度
        seq_len = input_ids.shape[1]
        # 选择前半部分作为选定的输入
        chosen_ids = input_ids[:bs]  # bs x seq x 1
        # 选择后半部分作为拒绝的输入
        rejected_ids = input_ids[bs:]
        # 前半部分选的的奖励
        chosen_rewards = rewards[:bs]
        # 后半部分被拒绝的奖励
        rejected_rewards = rewards[bs:]

        # Compute pairwise loss. Only backprop on the different tokens before padding
        # 计算配对损失，应用与强化学习中，主要目的是比较生成模型横城模型的不同序列（被接受和被拒绝的序列）的奖励
        loss = 0.
        for i in range(bs):
            # 遍历每个序列对（被接受和被拒绝的序列），提取相应的ID和奖励
            chosen_id = chosen_ids[i]
            rejected_id = rejected_ids[i]
            chosen_reward = chosen_rewards[i]
            rejected_reward = rejected_rewards[i]
            '''##################################################################################################'''
            # ! 这一部分代码的目的：获取被选定的序列和被拒绝的序列中差异性部分。也可以理解为回答的内容。同时，内容的长度是根据序列（去除padding）最长的那个决定。
            # 找出在被选定序列中那些位置存在填充（PAD）的标记，并选择第一个非填充标记的位置作为序列的结束结果
            # 如果不存在填充标记，则使用整个序列的长度作为结束位置
            c_inds = (chosen_id == self.PAD_ID).nonzero() # Pytorch中的.nonzero()方法，这个方法会返回输入张量中所有非零元素的索引
            c_ind = c_inds[self.num_padding_at_beginning].item() if len(
                c_inds
            ) > self.num_padding_at_beginning else seq_len  # OPT model pads the first token, so we need to use the second padding token as the end of the sequence
            # 查找被选定序列和被拒绝序列之间的差异位置，这里的差异位置是指这两个序列第一个不同的标记的位置  
            # ? 为什么要这样做
            check_divergence = (chosen_id != rejected_id).nonzero()
            # 如果两个序列完全相同，则结束位置设置为序列的最后一位，差异位置设置为结束位置的前一位
            if len(check_divergence) == 0:
                end_ind = rejected_reward.size(-1) # 获取序列长度
                divergence_ind = end_ind - 1 # 原因：虽然两个序列没有差异，但是我们仍然需要一个“差异点”，因为这个位置被用于计算奖励，所以选择最后一个
                # r_ind = c_ind 在这种没有差异的情况下，被选定的序列和被拒绝的序列应该有相同的长度
            else:
                # Check if there is any padding otherwise take length of sequence
                # 如果两个序列不完全相同，再次对被拒绝的序列执行相同的填充检查过程，并确定新的检查位置和差异位置
                r_inds = (rejected_id == self.PAD_ID).nonzero() # 找出所有的填充位置
                r_ind = r_inds[self.num_padding_at_beginning].item(
                ) if len(r_inds) > self.num_padding_at_beginning else seq_len # 确定结束位置end_ind,结束位置应该是两个序列的最大值
                end_ind = max(c_ind, r_ind)
                divergence_ind = check_divergence[0] # 确定差异位置
            assert divergence_ind > 0
            # 基于差异位置和结束位置，截取被选定序列和被拒绝序列的奖励部分
            c_truncated_reward = chosen_reward[divergence_ind:end_ind]
            r_truncated_reward = rejected_reward[divergence_ind:end_ind]
            '''##################################################################################################'''
            # 从截取的奖励中计算平均值并添加到被选定和被拒绝的平均列表中
            # 选择最后一个非填充token的奖励值添加到得分列表，可能是因为最后一个token往往包含了模型对整个输入序列的最终评价
            chosen_mean_scores.append(
                chosen_reward[c_ind - 1])  #use the end score for reference
            rejected_mean_scores.append(rejected_reward[r_ind - 1])

            if self.compute_fp32_loss:
                c_truncated_reward = c_truncated_reward.float()
                r_truncated_reward = r_truncated_reward.float()
            # 使用logsigmoid函数计算两个截取奖励之差的对数值，并将结果添加到损失中
            # logsigmoid是一种将任何实数映射到(0,1)区间的函数
            # 当c_truncated_reward大于r_truncated_reward时，这个值会接近1
            # 否则，它会接近于0，这样可以使的损失更加稳定，避免在梯度下降时，步长过大
            loss += -torch.nn.functional.logsigmoid(c_truncated_reward -
                                                    r_truncated_reward).mean()
        # 将模型在整个batch中每个样本上计算得到的损失值进行平均，得到一个标量值最为一个batch的损失 
        loss = loss / bs
        # 将两个分数列表堆叠为tensor
        chosen_mean_scores = torch.stack(chosen_mean_scores)
        # 返回包含损失值以及被选定和被拒绝序列的平均分数的字典 
        rejected_mean_scores = torch.stack(rejected_mean_scores)
        return {
            "loss": loss,
            "chosen_mean_scores": chosen_mean_scores,
            "rejected_mean_scores": rejected_mean_scores,
        }
    # ? 为什么这个没有涉及到 self.num_padding_at_beginning
    def forward_value(self,
                      input_ids=None,
                      attention_mask=None,
                      past_key_values=None,
                      position_ids=None,
                      head_mask=None,
                      inputs_embeds=None,
                      return_value_only=False,
                      prompt_length=0,
                      use_cache=False):

        if self.config.model_type == "llama":
            kwargs = dict()
        else:
            kwargs = dict(head_mask=head_mask)
        # 模型的隐藏状态，他是一个列表，包含了每一层的输出，每一层的输出都是一个形状为[batch_size, sequence]
        transformer_outputs = self.rwtransformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            **kwargs)
        # 取出第一层的隐藏状态
        hidden_states = transformer_outputs[0]
        # 这个值可以被视为模型对每个输入token的奖励
        values = self.v_head(hidden_states).squeeze(-1)
        if return_value_only:
            return values
        else:
            # [0 0 0 0 prompt, answer, 0 0 0 0 ] for step 3, we have padding at the beginning
            # [prompt, answer, 0, 0, 0, 0] this is normal
            # 序列的开始部分通常包含一些padding，如果问题部分长度不超过1，则无法准确获取答案部分的值
            assert prompt_length > 1, "prompt_length must be greater than 1 to help select the end score"
            bs = values.size(0)
            seq_len = input_ids.shape[1]
            # 保存每个输入序列答案部分的最后一个值
            chosen_end_scores = [
            ]  # we use this name for consistency with the original forward function
            for i in range(bs):
                # 获取当前序列的输入ID和模型计算出的值
                input_id = input_ids[i]
                value = values[i]
                # 查找当前序列答案部分（即问题部分之后的部分由input_id[prompt_length:]得到）中的padding标记
                # 这一步的结果是一个包含所有的padding标记位置的列表
                c_inds = (input_id[prompt_length:] == self.PAD_ID).nonzero()
                # here we only use the answer part of the sequence so we do not need to care about the padding at the beginning
                # 计算答案部分的结束位置，如果答案部分中存在padding标记（即len(c_inds) > 0）,
                # 结束位置就是第一个padding标记的位置；否则，结束位置就是整个序列的长度
                # 这里通过c_inds[0].item() + prompt_length来获取绝对位置，因为c_inds[0].item()只是相对于答案部分的开始位置的相对位置
                c_ind = c_inds[0].item() + prompt_length if len(
                    c_inds) > 0 else seq_len
                # 将答案部分的最后一个值（即结束位置前一个位置的值，由value[c_ind - 1]获取）加入到chosen_end_scores列表中
                chosen_end_scores.append(value[c_ind - 1])
            # 得到一个包含每个序列答案部分最后一个值的列表，这个列表可以用来表示每个序列的答案的“值”
            return {
                "values": values,
                "chosen_end_scores": torch.stack(chosen_end_scores),
            }
